import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from PIL import Image

import seaborn as sns

plt.style.use('ggplot')

"""
CONFIG DATA
"""

include_post = False
group = 2
BACKBONE = 'SERESNEXT50'

base1 = "/kaggle/input/xview2-assess-building-damage/train_images_labels_targets/train_images_labels_targets/train/"
base2 = "/kaggle/input/xview2-assess-building-damage/tier3/tier3/"



def read_meta(base_path, name_group=1):
    l_images = os.listdir(base_path+"images")
    summarize = {}
    summarize["collection"], summarize["id"], summarize["state"] = list(zip(*[el.split('_')[:-1] for el in l_images]))

    df = pd.DataFrame(summarize)
    df["group"] = name_group

    df.sort_values(by=["collection","id","state"], inplace=True)
    df = df.groupby(['collection', 'id','group'])['state'].agg('_'.join).reset_index()

    return df



#df = read_meta(base1, name_group=1)
#df2 = read_meta(base2, name_group=2)
#df = pd.concat([df1, df2])

df = pd.read_csv("/kaggle/input/building-segmentation-model-for-satellite-imagery/building_info.csv", dtype='str')

for col in ['group','no-damage','minor-damage','major-damage','destroyed']:
    df[col] = df[col].astype('int')


if group==1:
    # just train group 1
    df = df[df['group']==1]

# check the building count per image
build_count = df[['no-damage','minor-damage','major-damage','destroyed']].sum(axis=1)
df['b_count'] = build_count

df_yb = df[~(df[['no-damage','minor-damage','major-damage','destroyed']].sum(axis=1)==0)]
df_nb = df[(df[['no-damage','minor-damage','major-damage','destroyed']].sum(axis=1)==0)]

limit_group = 500

df_n = pd.DataFrame()
for name in df_nb['collection'].unique(): #groupby('collection')['id'].count()
    temp_df = df[df['collection']==name]
    
    # downsample no-building images
    if len(temp_df)>limit_group:
        temp_df = temp_df.sample(n=limit_group, replace=False, random_state=2024)

    df_n = pd.concat([df_n, temp_df], axis=0, ignore_index=True)

from sklearn.model_selection import train_test_split

train, test = train_test_split(df_n, random_state=2024, test_size=0.1) # 0.25


meta_tr = train.groupby("collection")['id'].count().to_dict()
meta_ts = test.groupby("collection")['id'].count().to_dict()

######### Ignore poorly defined label samples

ignore_l = ["socal-fire_00001233","hurricane-harvey_00000107",
            "hurricane-harvey_00000138","socal-fire_00001313",
            "hurricane-harvey_00000245","moore-tornado_00000142",
            "lower-puna-volcano_00000089"]

train = train[~(train['collection']+"_"+train['id']).isin(ignore_l)]

#########

#for k in meta_tr.keys():
#    print(f"{k:<20} {meta_tr[k]:>5} {meta_ts.get(k):}")# {meta_tr[k]/meta_ts[k]:>3}")

import tensorflow as tf
from tensorflow.keras import layers

tpu = None

try:
    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
    tf.config.experimental_connect_to_cluster(tpu)
    tf.tpu.experimental.initialize_tpu_system(tpu)
    #`tf.distribute.experimental.TPUStrategy` is deprecated
    strategy = tf.distribute.TPUStrategy(tpu)
except:
    strategy = tf.distribute.get_strategy()

!pip install shapely
!!pip install -U segmentation-models # automatic U-Net implementation for tf.keras
#!pip install -U image-classifiers # collection of novel classification models (2019), already installed with segmentation-models

from PIL import Image, ImageDraw
from shapely import wkt
import numpy as np
import json


# create binary mask from label json file
def create_build_from_json(file_path):

    data = tf.io.read_file(file_path)
    data = json.loads(data.numpy().decode('utf-8'))

    img = Image.new('L', (1024, 1024), color=0) #'black')
    draw = ImageDraw.Draw(img, 'L')
    #color = (255) #, 255, 255, 125) # white

    for cnt, polygon in enumerate( data['features']['xy'] ):
        x,y = wkt.loads(polygon['wkt']).exterior.coords.xy
        coords = list(zip(x,y))
        draw.polygon(coords, fill=1)

    return np.array(img)[...,np.newaxis] # dtype uint8

def crop_image(input_size, crop_size, stride):
    """
        Crop images from 1024x1024 to 512x512 patches from image
        inputs shape: (rows. columns. channels) -> output (#splits, rows. columns. channels)
    """

    def inner_funct(image, label):
        patches = []
        for i in range(0, input_size, stride):
            for j in range(0, input_size, stride):

                # patches inside image only
                if (i+crop_size)>input_size or (j+crop_size)>input_size:
                    continue

                patch_im = tf.image.crop_to_bounding_box(image, i, j, crop_size, crop_size)
                patch_lb = tf.image.crop_to_bounding_box(label, i, j, crop_size, crop_size)

                patches.append( (patch_im, patch_lb) )

        return tf.data.experimental.from_list(patches)
    return inner_funct

from tensorflow.keras.layers import Layer


class BinarizeMaskLayer(Layer):

    def __init__(self, threshold=0.5, **kwargs):
        super(BinarizeMaskLayer, self).__init__(**kwargs)
        self.threshold = threshold

    def call(self, inputs):
        # Apply the binarization threshold across the batch
        return tf.where(inputs > self.threshold, 1.0, 0.0)

    def get_config(self):
        config = super().get_config()
        config.update({"threshold": self.threshold})
        return config

# simple random flip applied to whole batch
class Augment(Layer):
    def __init__(self, seed=42):
        super().__init__()
        # both use the same seed, so they'll make the same random changes.
        # **MUST TWO** times cause, each call can generate differente effects

        self.aug_inputs = tf.keras.Sequential([
            tf.keras.layers.RandomRotation(factor=0.2, fill_mode="constant", seed=seed),
            tf.keras.layers.RandomFlip(mode='horizontal_and_vertical', seed=seed),
            tf.keras.layers.RandomTranslation(height_factor=0.1, width_factor=0.1, fill_mode="constant", seed=seed),
            tf.keras.layers.RandomZoom(height_factor=0.2, width_factor=0.2, fill_mode='constant', seed=seed),
            tf.keras.layers.RandomBrightness(factor=0.1, value_range=(0.0, 1.0), seed=seed),
            tf.keras.layers.RandomContrast(factor=0.1, seed=seed)

        ])

        self.aug_labels = tf.keras.Sequential([
            tf.keras.layers.RandomRotation(factor=0.2, fill_mode="constant", seed=seed),
            tf.keras.layers.RandomFlip(mode='horizontal_and_vertical', seed=seed),
            tf.keras.layers.RandomTranslation(height_factor=0.1, width_factor=0.1, fill_mode="constant", seed=seed),
            tf.keras.layers.RandomZoom(height_factor=0.2, width_factor=0.2, fill_mode='constant', seed=seed),
            BinarizeMaskLayer(threshold=0.5) # to avoid loss precision in border resolution
        ])



    def call(self, inputs, labels):
        inputs = self.aug_inputs(inputs)
        labels = self.aug_labels(labels)

        return inputs, labels

def has_nonzero_label(image, label):

    # Check if there are any non-zero elements in the label
    limit = int(512*512 * 0.025) # at leat 2.5% of building pixels per image
    mask_sum = tf.reduce_sum(tf.cast(tf.not_equal(label, 0), dtype="int32"))

    return mask_sum >= limit
    #return tf.reduce_any(tf.not_equal(label, 0))

from tensorflow.keras.layers import Conv2D, Reshape, Multiply, Permute, Activation, Dot, Add, MaxPooling2D, UpSampling2D
from tensorflow.keras import layers


class AttentionGate(tf.keras.layers.Layer):
    def __init__(self, filters, name):
        super(AttentionGate, self).__init__(name=f"AttentionGate_{name}")

        self.sigmoid = Activation('sigmoid')
        self.relu = layers.ReLU()

        # Define the convolutions
        self.conv_o = Conv2D(filters, kernel_size=1) # representation of the current pixel.
        self.conv_s = Conv2D(filters, kernel_size=1) # Representation of the pixels you’re attending to.
        self.psi    = Conv2D(1, kernel_size=1)

    def call(self, inputs, skip):

        out_a  = self.conv_o(inputs) # [batch_size, height, width, filters // 2]
        skip_a = self.conv_s(skip)   # [batch_size, height, width, filters // 2]
        psi_a  = self.psi(self.relu(out_a + skip_a)) # [batch_size, height, width, 1]
        attention = self.sigmoid(psi_a) # [batch_size, height, width, 1]
        skip = skip * attention

        return skip

from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Concatenate,
                                     UpSampling2D, Add, Activation, BatchNormalization)


def asymmetric_conv_block(input_tensor, filters, n_sx=''):
    # 1x3 convolution
    conv1 = Conv2D(filters, (1, 3), padding='same', name=n_sx+'_cv1')(input_tensor)
    # 3x1 convolution
    conv2 = Conv2D(filters, (3, 1), padding='same', name=n_sx+'_cv2')(input_tensor)
    # 3x3 convolution
    conv3 = Conv2D(filters, (3, 3), padding='same', name=n_sx+'_cv3')(input_tensor)

    # Combine the outputs of the asymmetric convolutions
    x = Add()([conv1, conv2, conv3])
    x = BatchNormalization(name=n_sx+'_bn')(x)
    x = Activation('relu')(x)

    return x


def decode_block(filters, x, concat_l, n_sx='', include_post_conv=True, include_attention_gate=False):
    x = Conv2DTranspose(filters, (3, 3), strides=(2, 2), padding='same', activation='swish', name=n_sx+'_conv_trans')(x)

    if len(concat_l)!=0:
        if include_attention_gate:
            skip_out = AttentionGate(filters=filters, name=n_sx)(x, concat_l[0])
            x = Concatenate()([x, skip_out])
        else:
            x = Concatenate()([x]+concat_l) #[el for el in concat_l])

    if include_post_conv:
        x = asymmetric_conv_block(x, filters=filters, n_sx=n_sx+'_acb')
    return x


def generate_model(enc_model, enc_layers):

    if enc_layers[0]==-1:
        enc5, enc4, enc3, enc2, enc1 = [enc_model.layers[name].output for name in enc_layers ]
    else:
        enc5, enc4, enc3, enc2, enc1 = [enc_model.get_layer(name=name).output for name in enc_layers]

    dec1 = decode_block(filters=512, x=enc5, concat_l=[enc4], n_sx='dec1', include_attention_gate=False)
    #print(dec1.shape)

    dec2 = decode_block(filters=256, x=dec1, concat_l=[enc3], n_sx='dec2', include_attention_gate=False)
    #print(dec2.shape)

    dec3 = decode_block(filters=128, x=dec2, concat_l=[enc2], n_sx='dec3', include_attention_gate=False)
    #print(dec3.shape)

    ## extra conections
    ## kernel should be twice the size of strides to be similar to "bilinear"
    ##extra1 = Conv2DTranspose(filters=64, kernel_size=(16, 16), strides=(8, 8), padding='same')(dec1)
    ##extra2 = Conv2DTranspose(filters=64, kernel_size=(8, 8), strides=(4, 4), padding='same')(dec2)

    #extra1 = UpSampling2D(size=(8,8), interpolation="bilinear")(dec1)
    #extra1 = Conv2D(64, kernel_size=(3,3), padding='same',activation="relu", name='ext1_cv')(extra1)

    #extra2 = UpSampling2D(size=(4,4), interpolation="bilinear")(dec2)
    #extra2 = Conv2D(64, kernel_size=(3,3), padding='same',activation="relu", name='ext2_cv')(extra2)

    dec4 = decode_block(filters=64, x=dec3, concat_l=[enc1], #extra1, extra2],
                        n_sx='dec4', include_attention_gate=False)

    # Output layer
    x = Conv2DTranspose(16, (3, 3), strides=(2, 2), padding='same')(dec4)
    #print(x.shape)
    outputs = Conv2D(1, (1, 1), activation='sigmoid', dtype="float32")(x)
    #print(outputs.shape)
    # Create model

    return tf.keras.Model(inputs=enc_model.input, outputs=outputs)

# Obtain alpha

"""
total = 0
positive_c = 0
for i, (t1_image, (t1_label, _)) in tqdm(enumerate(ds_train)):
    # to use float16 labels
    positive_c += int(tf.reduce_sum(tf.cast(t1_label, dtype="float32")).numpy())
    total      += tf.size(t1_label).numpy()

alpha = np.round( (total - positive_c)/total, 2)
print("alpha ", alpha)

"""

#alpha = 0.88 # train +tier3 # no empty slices

alpha = 0.9

from segmentation_models.base import Metric
import segmentation_models.base.functional as F



SMOOTH = 1e-5
class IOUScore(Metric):
    r""" The `Jaccard index`_, also known as Intersection over Union and the Jaccard similarity coefficient
    (originally coined coefficient de communauté by Paul Jaccard), is a statistic used for comparing the
    similarity and diversity of sample sets. The Jaccard coefficient measures similarity between finite sample sets,
    and is defined as the size of the intersection divided by the size of the union of the sample sets:

    .. math:: J(A, B) = \frac{A \cap B}{A \cup B}

    """

    def __init__(
            self,class_weights=None,class_indexes=None,threshold=0.5,
            per_image=False,smooth=SMOOTH,name=None,):

        name = name or 'iou_score'
        super().__init__(name=name)
        self.class_weights = class_weights if class_weights is not None else 1
        self.class_indexes = class_indexes
        self.threshold = threshold
        self.per_image = per_image
        self.smooth = smooth

    def __call__(self, gt, pr):
        # Explicitly cast inputs to float32 for consistency
        gt = tf.cast(gt, tf.float32)
        pr = tf.cast(pr, tf.float32)

        return F.iou_score(gt,pr,
            class_weights=self.class_weights,class_indexes=self.class_indexes,
            smooth=self.smooth,
            per_image=self.per_image, threshold=self.threshold, **self.submodules
        )



class FScore(Metric):
    r"""The F-score (Dice coefficient) can be interpreted as a weighted average of the precision and recall,
    where an F-score reaches its best value at 1 and worst score at 0.
    The relative contribution of ``precision`` and ``recall`` to the F1-score are equal.
    The formula for the F score is:

    .. math:: F_\beta(precision, recall) = (1 + \beta^2) \frac{precision \cdot recall}
        {\beta^2 \cdot precision + recall}
    """

    def __init__(self, beta=1, class_weights=None, class_indexes=None,
                 threshold=None, per_image=False, smooth=SMOOTH,name=None,):
        name = name or 'f{}-score'.format(beta)
        super().__init__(name=name)
        
        self.beta = beta
        self.class_weights = class_weights if class_weights is not None else 1
        self.class_indexes = class_indexes
        self.threshold = threshold
        self.per_image = per_image
        self.smooth = smooth

    def __call__(self, gt, pr):
        # Explicitly cast inputs to float32 for consistency
        gt = tf.cast(gt, tf.float32)
        pr = tf.cast(pr, tf.float32)

        return F.f_score(gt,pr,
            beta=self.beta,
            class_weights=self.class_weights,
            class_indexes=self.class_indexes,
            smooth=self.smooth,
            per_image=self.per_image,
            threshold=self.threshold,
            **self.submodules
        )

class CombLoss:
    def __init__(self, from_logits=False):
        self.from_logits = from_logits
        self.dice_loss = tf.keras.losses.Dice()
        self.binary_focal_loss = tf.keras.losses.BinaryFocalCrossentropy()

    def __call__(self, y_true, y_pred):
        loss = self.dice_loss(y_true, y_pred) + 10.0 * self.binary_focal_loss(y_true, y_pred)
        return loss

from IPython.display import clear_output

selected_slice = 27

def display(display_list):

    plt.figure(figsize=(15, 15))
    title = ['Input Image', 'True Mask', 'Predicted Mask']

    for i in range(len(display_list)):
        plt.subplot(1, len(display_list), i+1)
        plt.title(title[i])
        
        if i>=1:
            plt.imshow(display_list[i], vmin=0, vmax=1, cmap="gray")
        else:
            plt.imshow(display_list[i])
        plt.axis('off')

    plt.show()


def show_predictions():
    display([t_image[selected_slice],
             t_label[selected_slice],
             model.predict(t_image)[selected_slice].astype('float')])


class DisplayCallback(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        clear_output(wait=True)
        show_predictions()
        print('\nSample Prediction after epoch {}\n'.format(epoch+1))

#from transformers.modeling_tf_utils import keras

early_stopping = tf.keras.callbacks.EarlyStopping(monitor='val_loss', #'val_iou_score',
                                                  patience=4, mode='min', verbose=1,
                                                  min_delta=0, restore_best_weights=True)

# To reduce Python overhead and maximize the performance of your TPU, Produce erros with generator dataset
# pass in the steps_per_execution argument to Keras Model.compile.
#steps_per_epoch  = len(train_names[0]) // (GLOBAL_BATCH_SIZE*4)
#validation_steps = len(test_names[0]) // (GLOBAL_BATCH_SIZE*4)

history = model.fit(x=ds_train, epochs=128,
                    #steps_per_epoch=steps_per_epoch,
                    validation_data=ds_test, #validation_steps=validation_steps,
                    callbacks=[DisplayCallback(),
                               early_stopping])

from IPython.display import FileLink
# Save the entire model to a HDF5 file.

with strategy.scope():
    model.save_weights(f'bseg_unet_{BACKBONE}_v38_tpu.weights.h5')

FileLink(f'bseg_unet_{BACKBONE}_v38_tpu.weights.h5')

from tensorflow.keras.layers import Average


class Rot90Layer(tf.keras.layers.Layer):

    """
    A custom layer to rotate images by 90 degrees.

    Args:
        rotation (int): Number of 90-degree rotations. Positive values rotate counterclockwise.
    """

    def call(self, inputs, rotations=1):
        return tf.image.rot90(inputs, k=rotations)


def post_model(model):

    image = tf.keras.Input(shape=(512,512,3))
    rot_images = [Rot90Layer()(image, rotations=i) for i in range(1, 4)]
    pred_l = [model(image)] + [model(rot_image) for rot_image in rot_images]
    pred_l = [Rot90Layer()(pred, rotations=-1*i) for i, pred in enumerate(pred_l)]
    output = Average(dtype="float32")(pred_l)

    return tf.keras.Model(inputs=image, outputs=output)


with strategy.scope():
    ### base model
    enc_layers =  (-1, 1078, 584, 254, 4)
    classifier, preprocess_input = Classifiers.get('seresnext50') # 'seresnext50': (1078, 584, 254, 4)
    encoder_model = classifier(input_shape=(512,512,3), weights=None,include_top=False)
    encoder_model.trainable=True

    model_base = generate_model(enc_model=encoder_model, enc_layers=enc_layers)
    model_base.load_weights("/kaggle/input/building-segmentation-model-for-satellite-imagery/bseg_unet_SERESNEXT50_v38_tpu.weights.h5")
    
    #model_base = model
    
    ########## post model
    post_m = post_model(model_base)
    post_m.trainable=False
    
    post_m.compile(
        tf.keras.optimizers.AdamW(learning_rate=0.00015),
        loss= lambda y_true, y_pred: 0.0, # Dummy Loss
        metrics=[IOUScore(threshold=0.4), # Low considering that noise was reduced as much as possible
                 FScore(beta=0.5)],
    )

base2 = "/kaggle/input/xview2-assess-building-damage/test_images_labels_targets/test_images_labels_targets/test/"
images_test = tf.io.gfile.listdir(base2+"images")
images_test = [el.split('_pre')[0] for el in images_test if len(el.split('_pre'))==2]

############ ignore poorly defined labels
m_l = ['socal-fire_00001378', 'hurricane-harvey_00000133', 'socal-fire_00001029']
images_test = [el for el in images_test if el not in m_l]
############

images  = [base2+'images/' + el + '_pre_disaster.png' for el in images_test]
targets = [base2+'targets/'+ el + '_pre_disaster_target.png'for el in images_test]


ds_test2 = tf.data.Dataset.from_tensor_slices((images, targets))
ds_test2 = (ds_test2.map(preprocess_load_file(use_json=False),
                         num_parallel_calls=tf.data.AUTOTUNE)
            .flat_map(lambda image, label: (crop_function_test(image, label)))
            .batch(GLOBAL_BATCH_SIZE)
            .prefetch(tf.data.AUTOTUNE))

ds_test2 = tf.data.Dataset.from_tensor_slices((images, targets))
ds_test2 = (ds_test2.map(preprocess_load_file(flag=False, use_json=False), #cont=[mean, std]),
                         num_parallel_calls=tf.data.AUTOTUNE)
            #.flat_map(lambda image, label: (crop_function_test(image, label)))
            .batch(GLOBAL_BATCH_SIZE)
            .prefetch(tf.data.AUTOTUNE))

t2_image, t2_label = ds_test2.take(1).as_numpy_iterator().next()
t2_pred = model.predict(t2_image)